package com.example.djlspringbootstartdemo;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.ImageFactory;
import ai.djl.ndarray.types.DataType;
import ai.djl.ndarray.types.Shape;
import ai.djl.translate.Transform;
import ai.djl.translate.Translator;
import ai.djl.translate.TranslatorContext;
import ai.djl.ndarray.NDArray;
import ai.djl.ndarray.NDList;

/**
 * Created by lilijava
 * on 2024/12/23
 * <p>
 * on  djl-spring-boot-start-demo
 */
public class ConvertToBinary implements Transform {

    private float threshold = 0.5f; // 可以根据需要调整阈值


    public NDList transform(NDList src, NDList dest, TranslatorContext ctx) {
        // 假设src中只有一个NDArray，即图像数据
        NDArray array = src.head();
        // 将图像数据转换为灰度格式
        NDArray grayArray = array.toType(DataType.UINT8, true).mean(new int[]{2}, true);
        // 将单通道灰度图像扩展为三通道，以便与原始图像形状一致
        grayArray = grayArray.expandDims(2).broadcast(new Shape(array.getShape().get(0), array.getShape().get(1), 3));
        return new NDList(grayArray);
    }


    @Override
    public NDArray transform(NDArray ndArray) {
        return null;
    }
}
